/*
 * JointData.h
 *
 *  Created on: Dec 29, 2023
 *      Author: zhian
 */

#ifndef _JOINTDATA_H_
#define _JOINTDATA_H_

/**
 * 可以包括多个模型的joint 数据
 * byChain 模式每个array含有整个model chain的轴数据, 比如在外部轴机器人chain里这样的array含有所有轴数据
 * byModel 模式每个array含有每个model的轴数据， 比如在外部轴机器人chain里每个array是机器人或者外部轴的轴信息
 */
#include <Eigen/Dense>
#include <iostream>


class JointData
{
public:
    /**
     * default by chains
     */
    JointData(int numArrays, bool isByChain=true)
    {
        data.resize(numArrays);
        m_byChain = isByChain;
    }

    JointData() : m_byChain(true) {}
    /**
     * only set for a certain model chain
     */
    void setJoints(int modelChainID, const std::vector<double>& jts)
    {
        if ( m_byChain == false && ids.size() != 0 ) {
            std::cerr << "请把已有Model数据转换成Chain数据再设新chain 轴数据" << std::endl;
        }
        int i = findIndex(modelChainID);
        if ( i == -1 ) {
            ids.push_back(modelChainID);
            data.emplace_back(jts);
        }
        else {
            ids[i] = modelChainID;
            data[i] = jts;
        }
        //TO_DO: forbid mixing by chain and by model data
        m_byChain = true;
    }

    bool getJoints(int id, std::vector<double>& jts) const
    {
        int i = findIndex(id);
        if ( i == -1 ) return false;
        else {
            jts = data[i];
            return true;
        }
        return true;
    }

    std::vector<double> getJoints(int id) const
    {
        std::vector<double> ret;
        int i = findIndex(id);
        if ( i == -1 ) return ret;
        else {
            ret = data[i];
            return ret;
        }
    }
    /**
     *
     */

    void UpdateChainID(int newID)
    {
        if(0 > ids.front()){
            ids.front() = newID;
        }else{
            ids.front() = newID;
        }
    }

    /*
     * set for all model chains，mcids给出对应的model chain ID
     * 注意这个会冲掉所有已有数据
     */
    void setAllJoints(const std::vector<std::vector<double>>& jts, const std::vector<int>& mcids)
    {
        m_byChain = true;
        ids = mcids;
        data = jts;
    }
    /**
     * 设立某个model的轴数据
     */
    void setJointsByModel(int modelID, const std::vector<double>& jts)
    {
        int i = findIndex(modelID);
        if ( i == -1 ) {
            ids.push_back(modelID);
            data.emplace_back(jts);
        }
        else {
            ids[i] = modelID;
            data[i] = jts;
        }
        //TO_DO: forbid mixing by chain and by model data
        m_byChain = false;
    }

    /**
     * use for set by model mode， mids提供对应的 modelID， 注意两种模式不能混合使用
     * 注意这个会冲掉所有已有数据
     */
    void setAllJointsByModel(const std::vector<std::vector<double>>& jts, const std::vector<int>& mids)
    {
        m_byChain = false;
        ids = mids;
        data = jts;
    }

    inline int findIndex(int id) const {
        int ret = -1;
        for ( size_t i = 0; i < ids.size(); i++)  {
            if ( id == ids[i]) {
                ret = i;
                break;
            }
        }
        return ret;
    }

    std::vector<std::vector<double>> data;
    std::vector<int> ids;
    bool isByChain() const { return m_byChain; }
    void merge(const JointData& jd) {
        for ( size_t i = 0; i < jd.data.size(); i++) {
            ids.push_back(jd.ids[i]);
            data.push_back(jd.data[i]);
        }
    }

    bool isTheSameType(const JointData& jd) {
        bool b = std::equal(jd.ids.begin(), jd.ids.end(), ids.begin(), ids.end());
        return b;
    }

    /**
     * by chain or by model
     */
    bool m_byChain;
};

class MJointData
{
public:
    MJointData()
    {
        pos = Eigen::Affine3d::Identity();
    }

    int chainID = -1;
    JointData data;
    Eigen::Affine3d pos;
    bool moved = false;

    void ClearJoint()
    {
        data.ids.clear();
        data.data.clear();
    }

    bool checkAny(const int &cid) {
        // if(cid < 0) return false;
        return std::any_of(data.ids.begin(), data.ids.end(), [cid](int jt) {
            return jt == cid;
        });
    }

    static bool checkAny(const std::vector<MJointData> &ms, const int &cid) {

        return std::any_of(ms.begin(), ms.end(), [cid](const MJointData &jt) {
            return jt.chainID == cid;
        });
    }

    static void Merge(std::vector<MJointData> &ms, const MJointData& mjd) {
        if(ms.empty()){
            ms = {mjd};
        }else{
            if(ms.front().chainID == -1 ){
                ms = {mjd};
            }else{
                ms.emplace_back(mjd);
            }
        }
    }

    void Merge(const MJointData& mjd)
    {
        data.merge(mjd.data);
        chainID = mjd.chainID;
    }

    void UpdateChainID(int newID)
    {
        chainID = newID;
        data.UpdateChainID(newID);
    }
};

#endif /* JOINTDATA_H_ */
